# sqls_embedding.py
# usage: python ./src/vector/sqls_embedding.py

import os, sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../'))
sys.path.insert(0, project_root)
import argparse
import json
from typing import List, Generator
from model.ai_loader import clientEmbedding
from config import project_config


def get_embeddings(texts, model=project_config.embedding_model):
    # 确保 texts 是列表
    input_texts = [texts] if isinstance(texts, str) else texts

    if not input_texts:
        return []

    all_embeddings = []
    batch_size = project_config.embedding_batch_size

    try:
        for i in range(0, len(input_texts), batch_size):
            batch = input_texts[i:i + batch_size]
            response = clientEmbedding.embeddings.create(
                input=batch,
                model=model
            )
            # 提取当前批次的嵌入向量
            batch_embeddings = [data.embedding for data in response.data]
            all_embeddings.extend(batch_embeddings)

        return all_embeddings

    except Exception as e:
        print(f"生成嵌入时出错: {e}")
        return None


def write_embedding_to_log(log_file: str, text: str, embedding: List[float], model: str, source_file: str):
    """
    将单个文本及其嵌入向量写入日志文件（JSON Lines 格式）。

    参数:
        log_file: 日志文件路径
        text: 原始文本
        embedding: 嵌入向量（浮点数列表）
        model: 使用的模型名称
        source_file: 来源 SQL 文件路径
    """
    log_entry = {
        "text": text,
        "model": model,
        "dimension": len(embedding),
        "embedding": embedding
    }
    try:
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(json.dumps(log_entry, ensure_ascii=False) + '\n')
        print(f"✅ 已将 '{source_file}' 的嵌入写入日志: {log_file}")
    except Exception as e:
        print(f"❌ 写入日志文件时出错 ({source_file}): {e}")


def read_sql_files(directory: str) -> Generator[str, None, None]:
    """
    遍历目录，查找所有 .json 文件，按顺序读取其内容并返回 (文件路径, 内容) 的生成器。

    参数:
        directory: 要扫描的目录路径

    Yields:
        (file_path, content) 元组
    """
    if not os.path.exists(directory):
        raise FileNotFoundError(f"目录不存在: {directory}")
    if not os.path.isdir(directory):
        raise NotADirectoryError(f"指定路径不是目录: {directory}")

    for root, _, files in os.walk(directory):
        for file in sorted(files):  # 排序确保顺序一致
            if file.lower().endswith('.json'):
                file_path = os.path.join(root, file)
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read().strip()
                        if content:  # 忽略空文件
                            yield file_path, content
                        else:
                            print(f"🟡 跳过空文件: {file_path}")
                except Exception as e:
                    print(f"❌ 无法读取文件 {file_path}: {e}")


def main():
    """
    主函数：从指定目录读取所有 .json 文件，生成嵌入并向量化，可选择写入日志。
    """
    parser = argparse.ArgumentParser(description="批量读取指定目录下的所有 .json 文件，生成嵌入向量并写入日志。")

    parser.add_argument(
        '--embed-dir',
        type=str,
        default=project_config.sql_output_path , # './sql_output',  # 增加默认值
        help='包含 .json 文件的目录路径，将递归扫描该目录下所有 .json 文件。'
    )

    parser.add_argument(
        '--log',
        type=str,
        default='tmp_embedding_waiting.txt',  # 增加默认值
        help='指定输出日志文件路径（JSON Lines 格式），用于保存所有嵌入结果。'
    )

    args = parser.parse_args()
    sql_dir = args.embed_dir
    log_file_path = args.log

    if not os.path.exists(sql_dir):
        print(f"❌ 目录不存在: {sql_dir}")
        sys.exit(1)

    print(f"🔍 开始扫描目录: {sql_dir}")
    sql_files = list(read_sql_files(sql_dir))

    if not sql_files:
        print("🟡 指定目录中未找到任何有效的 .json 文件。")
        return

    print(f"📄 找到 {len(sql_files)} 个 .json 文件，开始生成嵌入...")

    # 提取所有文本内容用于批量嵌入
    texts_to_embed = [content for _, content in sql_files]
    model_used = project_config.embedding_model

    embeddings = get_embeddings(texts_to_embed, model=model_used)

    if not embeddings or len(embeddings) != len(texts_to_embed):
        print("❌ 嵌入生成失败或数量不匹配。")
        sys.exit(1)

    print(f"✅ 成功生成 {len(embeddings)} 个嵌入向量！开始写入日志...")

    # 逐个写入日志
    for (file_path, content), embedding in zip(sql_files, embeddings):
        write_embedding_to_log(log_file_path, content, embedding, model_used, file_path)

    print(f"🎉 所有嵌入已成功写入日志文件: {log_file_path}")


# --- 程序入口 ---
if __name__ == '__main__':
    main()